///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file formic/cholesky.h
///
/// \brief   header file for cholesky decomposition
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef FORMIC_CHOLESKY_HEADER
#define FORMIC_CHOLESKY_HEADER

#include <vector>
#include <cassert>
#include <complex>
#include <cmath>
#include <algorithm>

namespace formic {

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \brief  Base clase for objects that hold the cholesky decomposition source
  ///         
  ///////////////////////////////////////////////////////////////////////////////////////////////////

  template <class T> class CholeskySource {

    public:

      // get the dimension of the matrix to be decomposed
      virtual int chol_size() const = 0;

      // returns the threshold below which we may truncate matrix elements
      virtual double chol_thresh() const = 0;

      // function to return the ith diagonal element of the matrix to be decomposed
      virtual double chol_diag(const int i) const = 0;

      // get a column of the matrix to be decomposed
      virtual void chol_get_col(const int j, std::vector<T> & col) const = 0;

  };

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \brief  Base clase for objects that hold the cholesky decomposition matrix L.
  ///         
  ///   A = L L^(hermitian conjugate)
  ///         
  ///////////////////////////////////////////////////////////////////////////////////////////////////

  template <class T> class CholeskyTarget {

    public:

      // do any preparations necessary for storing m columns of the decomposition matrix L
      virtual void chol_prepare(const int m) = 0;

      // store a column of the decomposition matrix L
      virtual void chol_put_col(const int j, const std::vector<T> & col) = 0;

  };

  ///////////////////////////////////////////////////////////////////////////////////////////////////
  /// \brief  Performs a Cholesky decomposition of the symmetric semi-definite matrix.
  ///         
  /// \param[in]      source   object holding the matrix to be decomposed
  /// \param[in,out]  target   object to hold the resulting decomposition
  ///
  ///////////////////////////////////////////////////////////////////////////////////////////////////

  template <class T> inline void cholesky_decomp(const formic::CholeskySource<T> & source,
                                                 formic::CholeskyTarget<T> & target) {

    // some local functions to perform some simple operations on either real or complex numbers
    struct ZZ {

      static void zero(double & x) { x = 0.0; }
      static void one(double & x) { x = 1.0; }
      static double conj(const double x) { return x; }
      static double real(const double x) { return x; }
      static double imag(const double x) { return 0; }

      static void zero(std::complex<double> & x) { x = std::complex<double>(0.0, 0.0); }
      static void one(std::complex<double> & x) { x = std::complex<double>(1.0, 0.0); }
      static std::complex<double> conj(const std::complex<double> & x) { return std::conj(x); }
      static double real(const std::complex<double> & x) { return x.real(); }
      static double imag(const std::complex<double> & x) { return x.imag(); }

    };

    // create the zero and one scalars
    T zero, one;
    ZZ::zero(zero);
    ZZ::one(one);

    // get the matrix dimension
    const int n = source.chol_size();

    // get the truncation threshold
    const double thresh = source.chol_thresh();
    assert( thresh >= 1.0e-11 );

    // get the diagonal elements of the matrix
    std::vector<std::pair<double,int> > ordered_diagonal(n);
    for (int i = 0; i < n; i++) {
      ordered_diagonal[i].first = -source.chol_diag(i);
      ordered_diagonal[i].second = i;
      assert( ordered_diagonal[i].first <= 0.0 );
    }

    // sort the row indicies in order of largest-to-smallest diagonal elements
    std::sort(ordered_diagonal.begin(), ordered_diagonal.end());

    // create a short name for an iterator
    typedef typename std::vector<std::pair<T,int> >::const_iterator iterator;

    // count the number of diagonal elements we will use (we ignore diagonal elements below the threshold)
    int m = 0;
    for ( ; m < ordered_diagonal.size() && std::abs(ordered_diagonal[m].first) > thresh; m++) {}

    // create a vector to store the columns of the cholesky matrix
    std::vector<std::vector<T> > LT(m);
    for (int j = 0; j < m; j++)
      LT[j].assign(n, zero);

    // create a vector to hold a column from the source matrix
    std::vector<T> source_col(n, zero);

    // prepare the target to hold m columns of the cholesky matrix
    target.chol_prepare(m);

    // iteratively generate the columns of the cholesky matrix
    for (int j = 0; j < m; j++) {

      // get the un-sorted index of the column
      const int J = ordered_diagonal[j].second;

      // get the column from the source matrix
      source.chol_get_col(J, source_col);

      // compute the diagonal entry of the new column
      assert( ZZ::real(source_col[J]) == -ordered_diagonal[j].first );
      assert( ZZ::imag(source_col[J]) == 0.0 );
      LT[j][J] = source_col[J] + one * thresh / 1000.0;
      for (int q = 0; q < j; q++)
        LT[j][J] -= LT[q][J] * ZZ::conj(LT[q][J]);
      assert( ZZ::real(LT[j][J]) > 0.0 );
      assert( std::abs(ZZ::imag(LT[j][J])) < thresh );
      LT[j][J] = std::sqrt(LT[j][J]);

      // compute the off-diagonal entries
      for (int i = j+1; i < n; i++) {
        const int I = ordered_diagonal[i].second;
        LT[j][I] = source_col[I];
        for (int q = 0; q < j; q++)
          LT[j][I] -= LT[q][I] * ZZ::conj(LT[q][J]);
        LT[j][I] /= LT[j][J];
      }

      // place the column in the target
      target.chol_put_col(j, LT[j]);

    }

  }

//  template<class T> void test_cholesky(const int n);

} // end namespace formic

#endif
